import sys
import os

from tensorflow.keras.applications.inception_resnet_v2 import InceptionResNetV2
from tensorflow.keras.mixed_precision import experimental as mixed_precision
import tensorflow as tf
import numpy as np
from 数据处理 import 数组加载, csv文件加载, csv文件加载_dataset

from 通用 import config

# 禁用gpu
# os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
batch_size = 4
# 加载数组格式的原始数据: 适合较小的 可以加载进内存的数据
# num_classes, x_train, y_train = 数组加载.load(config.训练数据)
# num_classes, offset, x_train, y_train = csv文件加载.load(config.训练数据, "纵向")
# shape = x_train[0].shape
train_dataset, validate_dataset, shape, classnum = csv文件加载_dataset.load(config.训练数据, batch_size, "纵向", "分类")

# 将数组列表的多维输入转为一维输入
# x_train = tf.stack(x_train)
# y_train = tf.stack(y_train)

# 构建模型  keras.applications模块提供了带有预训练权值的深度学习模型，这些模型可以用来进行预测、特征提取和微调 InceptionResNetV2比resNet准确率更高
#  include_top:是否带分类器 input_tensor:可自定义张量 weights:权重初始化 pooling:现在很多神经网络不再池化 (可能会丢失原始信息)
model = InceptionResNetV2(include_top=True,
                          weights=None,
                          input_tensor=None,
                          pooling=None,
                          input_shape=shape, classes=classnum)

# 控制台打印模型结构
model.summary()
# 绘制模型 windows环境下依赖可能会缺失
# keras.utils.plot_model(model, "mini_resnet.png", show_shapes=True)
# 编译模型
model.compile(
    run_eagerly=True,
    # Optimizer 优化器
    optimizer=tf.keras.optimizers.Adam(),
    # Loss function to minimize  损失 SparseCategoricalCrossentropy, MeanSquaredError, Huber
    loss=tf.keras.losses.SparseCategoricalCrossentropy(),
    # List of metrics to monitor 指标 分类问题:SparseCategoricalAccuracy, Accuracy,
    # 回归问题:MeanAbsoluteError 平均绝对值误差(越小越好) MeanRelativeError(normalizer=[[np.max(y_train)]]*batch_size) 平均相对误差(越小越好, 长度为每批大小, 最后一批大小不确定需要特殊处理, 未解决)
    metrics=[tf.keras.metrics.SparseCategoricalAccuracy()], )

# 预取数据
# train_dataset.prefetch(batch_size)
# 混合精度, 仅适用于n卡gpu, cpu机器使用混合精度反而可能会慢
# policy = mixed_precision.Policy('mixed_float16')
# mixed_precision.set_policy(policy)

# 训练模型   x:输入数据,可以是array或张量或list array其中数组认为是多个输入维度   y:目标数据,和x对应   batch_size:每轮中每一批的大小  epochs:轮次
# shuffle:混淆训练和验证数据, 不过该混洗只是每批单独混洗, 如果一批只包含一个种类,依然会过拟合
# validation_split:验证数据占比,在每个周期结束时，模型将迭代验证数据集并计算验证损失和验证指标
# callbacks:回调 TensorBoard:可视化
# 训练是为拟合神经网络: 训练集的x为输入, 经过网络直到输出层, 得到预测的y, 和实际的y相比较得到准确率, 对比上一批次的准确率(损失和指标), 以决定优化方向(优化器)
model.fit(x=train_dataset, validation_data=validate_dataset, epochs=20, callbacks=[
    tf.keras.callbacks.TensorBoard(
        log_dir=config.tensorBoard_path,
        histogram_freq=0,  # How often to log histogram visualizations
        embeddings_freq=0,  # How often to log embedding visualizations
        update_freq="epoch",
    )  # How often to write logs (default: once per epoch)
])

# 评估模型
scores = model.evaluate(validate_dataset)
print('loss:', scores[0])
print('acc:', scores[1])

# 保存模型
model.save(config.model_name_分类)
print('---保存模型为: ' + config.model_name_分类)

# 使用模型预测 见predict.py
# y = model.predict(x=x_train)
# print('y:', y)
sys.exit(1)
